package com.zzl.study.demo.controller;

import com.alibaba.cloud.ai.memory.redis.RedissonRedisChatMemoryRepository;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import static org.springframework.ai.chat.memory.ChatMemory.CONVERSATION_ID;

/**
 * 基于sse模式的MCP客户端Controller
 *
 * @author: zhangzl
 * @date: 2025/10/14 9:34
 * @version: 1.0
 */
@RestController
@RequestMapping("/mcp/sse")
public class SseMcpClientController {

    /**
     * 最大记忆消息数
     */
    private final static int MAX_MESSAGES = 1000;

    private final ChatClient chatClient;

    public SseMcpClientController(ChatClient.Builder chatClientBuilder, RedissonRedisChatMemoryRepository redisChatMemoryRepository, ToolCallbackProvider tools) {
        MessageWindowChatMemory messageWindowChatMemory = MessageWindowChatMemory.builder()
                .chatMemoryRepository(redisChatMemoryRepository)
                .maxMessages(MAX_MESSAGES)
                .build();
        this.chatClient = chatClientBuilder
                .defaultAdvisors(MessageChatMemoryAdvisor.builder(messageWindowChatMemory).build())
                .defaultToolCallbacks(tools)
                .build();
    }

    /**
     * 查询诗词
     * http请求头：Accept: text/event-stream
     *
     * @param input          用户输入
     * @param conversationId 会话ID
     * @return 诗词
     */
    @GetMapping("/query")
    public Flux<String> query(@RequestParam("input") String input, @RequestParam("conversation_id") String conversationId) {
        return this.chatClient
                .prompt(input)
                .advisors(
                        a -> a.param(CONVERSATION_ID, conversationId)
                )
                .stream()
                .content();
    }

}
